// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#pragma once

#include <cstdint>
#include <functional>
#include <map>
#include <memory>
#include <optional>
#include <queue>
#include <set>
#include <string>
#include <vector>

#include "coin/feed/fastfeed/feed.pb.h"
#include "coin1/base/filesystem.h"
#include "coin1/fastfeed/double_shift_feed_builder.h"
#include "orio/io/base.h"
#include "orio/record/simple.h"

namespace fastfeed {
namespace detail {

class ReaderContext {
 public:
  ReaderContext(
      const std::vector<std::experimental::filesystem::path>& filepaths,
      std::function<void(const proto::Feed&)> callback);
  virtual ~ReaderContext() {}

  virtual bool ReadNext();

  virtual void InvokeCallback() { callback_(feed_); }

  const std::vector<std::experimental::filesystem::path>& filepaths() const { return filepaths_; }
  virtual int64_t timestamp() const { return feed_.timestamp(); }
  virtual int64_t sort_key() const { return feed_.sort_key(); }

 protected:
  bool ReadNextAndParse(proto::Feed* feed);

  const std::vector<std::experimental::filesystem::path> filepaths_;
  std::function<void(const proto::Feed&)> callback_;
  proto::Feed feed_;
  std::unique_ptr<char[]> buf_;
  std::unique_ptr<orio::io::Reader<proto::Feed>> reader_;

  static constexpr size_t max_record_size = 16 * 1024 * 1024;
};  // class ReaderContext

class FrontRunReaderContext : public ReaderContext {
 public:
  FrontRunReaderContext(
      const std::vector<std::experimental::filesystem::path>& filepaths,
      std::function<void(const proto::Feed&)> callback);

  bool ReadNext() final;

  void InvokeCallback() final { callback_(*feed_); }

  int64_t timestamp() const final { return (*feed_).timestamp(); }
  int64_t sort_key() const final { return (*feed_).sort_key(); }

 private:
  std::array<proto::Feed, 2> feeds_;
  proto::Feed* feed_;
  proto::Feed* feed_next_;
  bool next_ready_ = false;
  bool result_ = false;
  bool result_next_ = false;
};  // class FrontRunReaderContext

class ReaderContextCompare {
 public:
  bool operator()(
      const std::shared_ptr<detail::ReaderContext> a,
      const std::shared_ptr<detail::ReaderContext> b) const {
    if (a->timestamp() != b->timestamp()) {
      return a->timestamp() > b->timestamp();
    } else {
      return a->sort_key() > b->sort_key();
    }
  }
};

}  // namespace detail

struct FeedCacheKey {
  std::string machine;
  std::string worker_id;
  std::string sub_req;

  bool operator==(const FeedCacheKey& other) const {
    return (machine == other.machine && worker_id == other.worker_id && sub_req == other.sub_req);
  }

  bool operator<(const FeedCacheKey& other) const {
    int r1 = machine.compare(other.machine);
    if (r1 != 0) {
      return r1 < 0;
    } else {
      int r2 = worker_id.compare(other.worker_id);
      if (r2 != 0) {
        return r2 < 0;
      } else {
        return sub_req < other.sub_req;
      }
    }
  }

  std::string ToString() const;
};  // struct FeedCacheKey

struct FeedCacheSubReq {
  FeedCacheKey feed_cache_key;
  std::vector<std::string> symbols;
};  // struct FeedCacheSubReq

class FeedCacheReader {
 public:
  FeedCacheReader(
      const std::string& date_str,
      const std::vector<filesystem::path>& feed_cache_dirs);
  ~FeedCacheReader() = default;

  void AddCacheKey(
      const FeedCacheKey& feed_cache_key,
      const std::vector<std::string>& symbols,
      std::function<void(const proto::Feed&)> callback);

  bool Read();
  std::optional<int64_t> GetNextTimestamp() const;

 private:
  std::vector<std::experimental::filesystem::path> ResolveCacheFilePath(
      const FeedCacheKey& cache_key,
      const std::string& symbol);

  void PushOrDeleteReader(std::shared_ptr<detail::ReaderContext> reader);
  std::shared_ptr<detail::ReaderContext> PopReader();

  std::priority_queue<
      std::shared_ptr<detail::ReaderContext>,
      std::vector<std::shared_ptr<detail::ReaderContext>>,
      detail::ReaderContextCompare>
      pqueue_;

  const std::string date_str_;
  const std::vector<filesystem::path> feed_cache_dirs_;
};  // class FeedCacheReader

class FeedCacheRunner {
 public:
  FeedCacheRunner(
      std::function<void(const FeedCacheKey&, FeedBuilderBase*)> feed_builder_reset_callback,
      const std::string& date_str,
      const std::vector<filesystem::path>& feed_cache_dirs);
  ~FeedCacheRunner() = default;

  void AddCacheKey(const FeedCacheKey& feed_cache_key, const std::vector<std::string>& symbols);
  void AddCacheKey(const FeedCacheSubReq& request);
  void AddCacheKeyWithFeedArb(
      const std::string& machine,
      const std::string& sub_req,
      const std::vector<std::string>& symbols);

  void Run();
  void Run(std::optional<int64_t> from_timestamp, std::optional<int64_t> to_timestamp);

 private:
  struct FeedBuilderHolder {
    std::set<std::string> symbols;
    std::set<std::string> worker_ids;
    std::unique_ptr<DoubleShiftFeedBuilder> feed_builder;
  };  // struct FeedBuilderHolder
  std::map<FeedCacheKey, FeedBuilderHolder> req_;

  std::function<void(const FeedCacheKey&, FeedBuilderBase*)> feed_builder_reset_callback_;
  const std::string date_str_;
  const std::vector<filesystem::path> feed_cache_dirs_;
};  // class FeedCacheRunner

}  // namespace fastfeed
